import ffmpeg
import pandas as pd
import os
import tempfile
from tqdm import tqdm

# 尝试导入ffmpeg-python，如果失败则提示安装正确的包
try:
    import ffmpeg
    getattr(ffmpeg, 'probe')
except (ImportError, AttributeError):
    raise ImportError("请安装ffmpeg-python包，而不是ffmpeg包。使用命令: pip install ffmpeg-python")

def get_video_info(video_path):
    """
    获取视频的详细信息
    
    :param video_path: 视频文件路径
    :return: 视频信息字典
    """
    try:
        probe = ffmpeg.probe(video_path)
        video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
        audio_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'audio'), None)
        format_info = probe.get('format', {})
        
        info = {
            'duration': float(format_info.get('duration', 0)),
            'video_stream': video_stream,
            'audio_stream': audio_stream
        }
        return info
    except Exception as e:
        print(f"获取视频信息时出错: {e}")
        return None

def get_video_frame_rate(video_path):
    """
    获取视频的帧率
    
    :param video_path: 视频文件路径
    :return: 视频帧率 (fps)
    """
    try:
        probe = ffmpeg.probe(video_path)
        video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
        if video_stream is None:
            raise ValueError("未找到视频流")
        
        # 获取平均帧率
        avg_frame_rate = video_stream['avg_frame_rate']
        if '/' in avg_frame_rate:
            num, den = map(int, avg_frame_rate.split('/'))
            if den != 0:
                return num / den
            else:
                # 如果分母为0，尝试使用r_frame_rate
                r_frame_rate = video_stream['r_frame_rate']
                if '/' in r_frame_rate:
                    num, den = map(int, r_frame_rate.split('/'))
                    if den != 0:
                        return num / den
        
        # 如果无法解析分数形式，尝试直接转换为浮点数
        try:
            return float(avg_frame_rate)
        except ValueError:
            pass
            
        # 默认返回30fps
        return 30.0
    except Exception as e:
        print(f"获取视频帧率时出错: {e}")
        # 默认返回30fps
        return 30.0

def extract_audio_segments(video_a_path, timestamps_csv, output_audio_path, frame_rate_b):
    """
    根据时间戳从视频A中提取音频片段并合并。

    :param video_a_path: 视频A的路径 (包含音频)
    :param timestamps_csv: 包含视频B帧对应视频A时间戳的CSV文件路径
    :param output_audio_path: 输出合并后音频的路径（.wav格式）
    :param frame_rate_b: 视频B的帧率 (fps)
    """
    # 获取视频信息
    video_info = get_video_info(video_a_path)
    if video_info is None:
        raise ValueError("无法获取视频信息")
    
    video_duration = video_info['duration']
    print(f"源视频总时长: {video_duration:.3f} 秒")
    
    # 1. 读取时间戳
    df_timestamps = pd.read_csv(timestamps_csv)
    timestamps = df_timestamps['start_time'].tolist()
    
    if not timestamps:
        raise ValueError("时间戳文件为空或格式不正确。")

    # 计算每帧的时长 (秒)
    duration_per_frame = 1.0 / frame_rate_b

    # 2. 提取音频片段
    temp_audio_files = []
    temp_dir = tempfile.mkdtemp()
    concat_file_path = None  # 初始化变量
    successful_segments = 0  # 记录成功提取的片段数

    try:
        for i, start_time in enumerate(tqdm(timestamps, desc="提取音频片段")):
            # 确定片段结束时间
            end_time = start_time + duration_per_frame 
            
            # 检查时间戳是否在有效范围内
            if start_time >= video_duration:
                print(f"警告: 片段 {i+1} 的开始时间 ({start_time:.3f}s) 超出视频时长 ({video_duration:.3f}s)，跳过处理")
                continue
                
            # 如果结束时间超出视频时长，则调整为视频结束时间
            if end_time > video_duration:
                end_time = video_duration
                print(f"提示: 片段 {i+1} 的结束时间调整为视频结束时间 {video_duration:.3f}s")

            # 定义临时文件名
            temp_file = os.path.join(temp_dir, f"segment_{i:06d}.wav")
            temp_audio_files.append(temp_file)

            # 使用 ffmpeg-python 提取音频片段
            try:
                (
                    ffmpeg
                    .input(video_a_path, ss=start_time, to=end_time) # ss: start time, to: end time
                    .output(temp_file, acodec='pcm_s16le', ac=1, ar=16000) # 转换为WAV格式以便合并 (可选)
                    .overwrite_output()
                    .run(quiet=True) # 静默运行以减少输出
                )
                # print(f"已提取片段 {i+1}/{len(timestamps)}: {start_time:.3f}s - {end_time:.3f}s")
                successful_segments += 1
            except ffmpeg.Error as e:
                print(f"处理第 {i+1} 个片段时出错:")
                print(f"开始时间: {start_time:.3f}s, 结束时间: {end_time:.3f}s")
                print(f"输入文件: {video_a_path}")
                print(f"输出文件: {temp_file}")
                # 继续处理下一个片段而不是中断整个过程
                continue

        print(f"成功提取 {successful_segments}/{len(timestamps)} 个音频片段")
        
        # 3. 检查是否有成功提取的音频片段
        if successful_segments == 0:
            raise ValueError("没有成功提取任何音频片段")

        # 创建一个包含所有片段文件名的文本文件供 ffmpeg 使用
        concat_file_path = os.path.join(temp_dir, "concat_list.txt")
        with open(concat_file_path, 'w') as f:
            for temp_file in temp_audio_files:
                # 只添加确实存在的文件
                if os.path.exists(temp_file):
                    # 使用 file 协议确保路径被正确处理
                    f.write(f"file '{os.path.abspath(temp_file)}'\n")

        # 检查concat文件是否为空
        if os.path.getsize(concat_file_path) == 0:
            raise ValueError("没有有效的音频片段用于合并")

        # 使用 ffmpeg 合并音频为WAV格式
        (
            ffmpeg
            .input(concat_file_path, format='concat', safe=0) # safe=0 允许非相对路径
            .output(output_audio_path, format='wav') # 输出为WAV格式
            .overwrite_output()
            .run(quiet=True)
        )
        print(f"音频合并完成: {output_audio_path}")

    finally:
        # 4. 清理临时文件
        for temp_file in temp_audio_files:
            if os.path.exists(temp_file):
                os.remove(temp_file)
        if concat_file_path and os.path.exists(concat_file_path):
            os.remove(concat_file_path)
        # 确保临时目录存在再尝试删除
        if os.path.exists(temp_dir):
            os.rmdir(temp_dir)

def merge_audio_video(video_b_path, new_audio_path, output_video_path):
    """
    将新音频合并到视频B中。

    :param video_b_path: 视频B的路径 (无声)
    :param new_audio_path: 新的音频文件路径（.wav格式）
    :param output_video_path: 最终输出视频的路径
    """
    # 检查输入文件是否存在
    if not os.path.exists(video_b_path):
        raise FileNotFoundError(f"视频文件不存在: {video_b_path}")
    
    if not os.path.exists(new_audio_path):
        raise FileNotFoundError(f"音频文件不存在: {new_audio_path}")

    # 使用 ffmpeg-python 合成视频和音频
    input_video = ffmpeg.input(video_b_path)
    input_audio = ffmpeg.input(new_audio_path)

    # 使用 'c:v copy' 和 'c:a aac' 可以无损复制视频流，同时将WAV转码为AAC
    try:
        ffmpeg.output(
            input_video['v'],  # 视频流
            input_audio['a'],  # 音频流
            output_video_path,
            vcodec='copy', # 复制视频流
            acodec='aac'   # 使用AAC编码音频 (将WAV转为AAC)
        ).overwrite_output().run(quiet=True) # 覆盖输出文件
        print(f"视频与音频合成完成: {output_video_path}")
    except ffmpeg.Error as e:
        print(f"视频音频合成时出错:")
        if e.stderr:
            print(f"ffmpeg stderr: {e.stderr.decode('utf-8')}")
        raise

# --- 主程序 ---
if __name__ == "__main__":
    # --- 请根据你的实际情况修改以下路径和参数 ---
    path_to_video_a = "6月13日.mp4"  # 包含音频的完整视频A
    path_to_video_b = "tmpjh9mrs4v.mp4"      # 从视频A剪辑出的无声视频B
    path_to_timestamps_csv = "video_timestamps_20250901_144719.csv"   # 视频B帧在视频A中的时间戳文件
    output_final_video = "result.mp4" # 输出的最终视频
    
    # 自动获取视频B的帧率
    frame_rate_of_video_b = get_video_frame_rate(path_to_video_b)
    print(f"检测到视频B的帧率为: {frame_rate_of_video_b} fps")

    # 临时文件路径 (使用.wav扩展名)
    temp_restored_audio = "temp_restored_audio.wav" # 临时音频文件

    try:
        print("开始提取并合并音频...")
        extract_audio_segments(path_to_video_a, path_to_timestamps_csv, temp_restored_audio, frame_rate_of_video_b)
        
        print("开始合成最终视频...")
        merge_audio_video(path_to_video_b, temp_restored_audio, output_final_video)
        
    except Exception as e:
        print(f"处理过程中发生错误: {e}")
    finally:
        # 清理临时音频文件
        if os.path.exists(temp_restored_audio):
            os.remove(temp_restored_audio)
            print(f"已清理临时文件: {temp_restored_audio}")

    print("所有处理完成。")